# -*- coding: utf-8 -*-

# File Name： norm
# Description :
# Author : lirui
# create_date： 2022/6/3
# Change Activity:

from torch import nn

from ....common.design_pattern.register.register_factory import REGISTER_FACTORY

_registry_name = 'norm_layer'


def _build_registry():
    """
    build normal lay registry.

    """
    _register = REGISTER_FACTORY.create_registry(_registry_name)
    _register.add('BN', value=nn.BatchNorm2d)
    _register.add('BN1d', value=nn.BatchNorm1d)
    _register.add('BN2d', value=nn.BatchNorm2d)
    _register.add('BN3d', value=nn.BatchNorm3d)
    _register.add('SyncBN', value=nn.SyncBatchNorm)
    _register.add('GN', value=nn.GroupNorm)
    _register.add('LayerNorm', value=nn.LayerNorm)
    _register.add('IN1d', value=nn.InstanceNorm1d)
    _register.add('IN2d', value=nn.InstanceNorm2d)
    _register.add('IN3d', value=nn.InstanceNorm3d)


_build_registry()  # build when module load.


def build_norm_layer(cfg=None, name=None):
    """
    build norm layer
    
    Args:
        name: module name
        cfg: must contain num_features;

    Returns:

    """

    return REGISTER_FACTORY.build_from_config(_registry_name, cfg, name)
